require(glmnet)
require(Matrix)
require(ggplot2)

demo <- function( n, K, p=.9 )
  {
    c <- round( n/K )
    n <- max( n, K*c )

    stopifnot( c > 2 ) # give up for too small group size

    A <- spMatrix( nrow=n, ncol=n )

    for( k in 1:K )
      {
        idx <- (1 + (k-1)*c):(k*c)
        A[idx,idx] <- runif( n=c*c )
      }

    A <- triu(A, k=1)
    A <- A + t(A)
    A[ A > p ] <- 0
    A[ A > 0 & A <= p ] <- 1
    return( A )
  }


################################################################
latent.seeding <- function( A, K )
  {

    rand.order <- sample(n,n)
    
    m <- min(n,2*K)
    clust <- 1:min(K,m)
    if( m > K ){ clust <- c(clust, 1:(m-K)) }
    
    ## random seeding for 1st K
    Y.rand <- spMatrix( nrow=n, ncol=K,
                       i = rand.order[1:m],
                       j = clust,
                       x = rep(1,m) )
    ## greedy seeding
    if( n > m ) {                    
      n.k <- apply(Y.rand, 2, sum)
      
      for( i in rand.order[-(1:m)] )
        {
          adj <- which( A[,i]>0 )
          if( length(adj) == 0 ){
            Y.rand[i,sample(K,1)] <- 1
          }
          if( length(adj) == 1 ){
            Y.rand[i,adj] <- 1
            next
          }
          d.k <- apply( Y.rand[adj,], 2, sum )
          if( sum(d.k) <= 0 ){ next }
          valid <- which(d.k>0)
          
          log.mass <- d.k[valid] - n.k[valid]
          k <- valid[which.max(log.mass)]
          Y.rand[i,k] <- 1
          n.k[k] <- n.k[k] + 1
        }
    }
    return( Y.rand )
  }

sigmoid <- function(x) 1/(1 + exp(-x))

################################################################
K <- 10
n <- 100
A <- demo( n, K=3, p=0.8 )


################################################################
## initial membership
K <- 15
X <- A + diag( rep(1,n) )

Y <- latent.seeding( A, K )

## fit initial beta 
valid <- apply(Y,2,sum) > 0
if( sum(!valid) > 0 ) {  beta[ ! valid , ] <- 0 }
out <- glmnet( X, as.matrix(Y[,valid]), family="multinomial", intercept=F, alpha=.1, lambda=.1 )
beta[valid,] <- t( sapply( out$beta, as.vector ) )


## re-estimate Y
argmax <- apply( A, 1, function(ai)
                which.max( sapply( 1:K, function(k) sum( ai * beta[k,] ) ) ) )
                
Y <- spMatrix( i=1:n, j=argmax, x=rep(1,n), nrow=n, ncol=K )

## re-estimate beta
valid <- apply(Y,2,sum) > 0
if( sum(!valid) > 0 ) {  beta[ ! valid , ] <- 0 }
out <- glmnet( X, as.matrix(Y[,valid]), family="multinomial", intercept=F, alpha=.1, lambda=.1 )
beta[valid,] <- t( sapply( out$beta, as.vector ) )


require(reshape2)
require(ggplot2)
p <- ggplot( melt( beta ), aes(x=Var2, y=Var1, fill=value) ) + geom_tile() +
  scale_fill_gradient2( low="red", mid="white", high="blue" ) + theme_bw()
print(p)





par(mfrow=c(2,2))
glmnet::plotCoef( out$beta[[1]], lambda=out$lambda, df=out$df, xvar="norm" )
glmnet::plotCoef( out$beta[[2]], lambda=out$lambda, df=out$df, xvar="norm" )
glmnet::plotCoef( out$beta[[3]], lambda=out$lambda, df=out$df, xvar="norm" )
